#define USE_CUDA

#include <iostream>
#include <iomanip>
#include "inference.h"
#include <filesystem>
#include <fstream>
#include <random>
#include "..//include/image_loading.hpp"

cv::Mat funtion(cv::Mat& img) {
    // 静态指针，只初始化一次
    static YOLO_V8* yoloDetector = nullptr;
    static bool isInit = false;
    if (!isInit) {
        yoloDetector = new YOLO_V8;
        yoloDetector->classes = { "0","1","2","3" };
        DL_INIT_PARAM params;
        params.rectConfidenceThreshold = 0.1;
        params.iouThreshold = 0.5;
        params.modelPath = "/home/nano/Documents/rm_f_03/yolo/yolov8s.onnx";
        params.imgSize = { 640, 640 };
    #ifdef USE_CUDA
        params.cudaEnable = true;
        params.modelType = YOLO_DETECT_V8;
    #else
        params.modelType = YOLO_DETECT_V8;
        params.cudaEnable = false;
    #endif
        yoloDetector->CreateSession(params);
        isInit = true;
    }

    std::vector<DL_RESULT> res;
    yoloDetector->RunSession(img, res);

    // 直接在原图上画框，减少 clone
    for (auto& re : res)
    {
        cv::RNG rng(cv::getTickCount());
        cv::Scalar color(rng.uniform(0, 256), rng.uniform(0, 256), rng.uniform(0, 256));
        cv::rectangle(img, re.box, color, 3);

        float confidence = floor(100 * re.confidence) / 100;
        std::string label = yoloDetector->classes[re.classId] + " " +
            std::to_string(confidence).substr(0, std::to_string(confidence).size() - 4);

        cv::rectangle(
            img,
            cv::Point(re.box.x, re.box.y - 25),
            cv::Point(re.box.x + label.length() * 15, re.box.y),
            color,
            cv::FILLED
        );

        cv::putText(
            img,
            label,
            cv::Point(re.box.x, re.box.y - 5),
            cv::FONT_HERSHEY_SIMPLEX,
            0.75,
            cv::Scalar(0, 0, 0),
            2
        );
    }
    return img;
}